#include "GDI.h"



NTSTATUS GDI::GDIInitialization()
{
	
	NTSTATUS Status = m_ShadowSSDT.InitializationShadowSSDT();
	if (!NT_SUCCESS(Status))
	{
		DPRINT("GDI.cpp InitializationShadowSSDT() An Error.GDIInitialization Internal Function\r\n");
		return Status;
	}

	Status = KernelGDIFunctionInitialization();
	if (!NT_SUCCESS(Status))
	{
		DPRINT("GDI.cpp KernelGDIFunctionInitialization() An Error.KernelGDIFunctionInitialization Internal Function\r\n");
		return Status;
	}

	return Status;

}

BOOLEAN GDI::GDIStartsCalling(HANDLE Pid, PKAPC_STATE Kpc, PULONG64  OldThreadWin32)
{
	if (Pid < 0)
	{
		DPRINT("GDI.cpp Pid < 0 An Error.GDIStartsCalling Internal Function\r\n");
		return FALSE;
	}

	NTSTATUS Status = Process::KeAttchProcess_Pfn(Pid, Kpc);
	if (NT_SUCCESS(Status))
	{
		PVOID Win32Thread = Thread::GetWin32ThreadAddress(IoGetCurrentProcess());
		PVOID Win32Proces = Process::PsGetCurrentProcessWin32Process();
		*OldThreadWin32 = (ULONG64)Thread::PsGetCurrentThreadWin32Thread();
		if (Win32Thread == NULL) {

			GDIStopsCalling(*Kpc, (ULONG64)Thread::PsGetCurrentThreadWin32Thread());
			DPRINT("GDI.cpp GetWin32ThreadAddress() An Error.GDIStartsCalling Internal Function\r\n");
			return FALSE;
		}

		if (Win32Proces == NULL) {


			GDIStopsCalling(*Kpc, (ULONG64)Thread::PsGetCurrentThreadWin32Thread());
			DPRINT("GDI.cpp PsGetCurrentProcessWin32Process() An Error.GDIStartsCalling Internal Function\r\n");
			return FALSE;
		}

		Thread::PsSetThreadWin32Thread(KeGetCurrentThread(), Win32Thread, Thread::PsGetCurrentThreadWin32Thread());
		return TRUE;
	}
	DPRINT("GDI.cpp KeAttchProcess_Pfn() An Error.GDIStartsCalling Internal Function\r\n");
	return FALSE;
}

VOID GDI::GDIStopsCalling(KAPC_STATE Kpc,ULONG64 OldWin32)
{
	Thread::PsSetThreadWin32Thread(KeGetCurrentThread(), (PVOID)OldWin32, Thread::PsGetCurrentThreadWin32Thread());
	Process::UnKeAttchProcess_Pfn(Kpc);
}

NTSTATUS GDI::KernelGDIFunctionInitialization()
{
	NTSTATUS Status = STATUS_UNSUCCESSFUL;
	m_KernelGDIFunction.NtUserFindWindowEx = (HWND(__cdecl *)(HWND, HWND, PUNICODE_STRING, PUNICODE_STRING, DWORD))m_ShadowSSDT.Get3SDTFunAddress(L"NtUserFindWindowEx");
	if (m_KernelGDIFunction.NtUserFindWindowEx == NULL || MmIsAddressValid(m_KernelGDIFunction.NtUserFindWindowEx) == FALSE) {
		DPRINT("GDI.cpp NtUserFindWindowEx() = NULL An Error.KernelGDIFunctionInitialization Internal Function\r\n");
		goto $EXIT;
	}

	m_KernelGDIFunction.NtGdiExtTextOutW = (BOOL(__cdecl *)(HDC, INT, INT, UINT, LPRECT, LPWSTR, INT, LPINT, DWORD))m_ShadowSSDT.Get3SDTFunAddress(L"NtGdiExtTextOutW");
	if (m_KernelGDIFunction.NtGdiExtTextOutW == NULL || MmIsAddressValid(m_KernelGDIFunction.NtGdiExtTextOutW) == FALSE) {
		DPRINT("GDI.cpp NtGdiExtTextOutW() = NULL An Error.KernelGDIFunctionInitialization Internal Function\r\n");
		goto $EXIT;
	}

	m_KernelGDIFunction.NtUserGetDC = (HDC(__cdecl *)(HWND))m_ShadowSSDT.Get3SDTFunAddress(L"NtUserGetDC");
	if (m_KernelGDIFunction.NtUserGetDC == NULL || MmIsAddressValid(m_KernelGDIFunction.NtUserGetDC) == FALSE) {
		DPRINT("GDI.cpp NtUserGetDC() = NULL An Error.KernelGDIFunctionInitialization Internal Function\r\n");
		goto $EXIT;
	}

	m_KernelGDIFunction.NtGdiPolyPolyDraw = (ULONG_PTR(__cdecl *)(HDC, PPOINT, PULONG, ULONG, INT))m_ShadowSSDT.Get3SDTFunAddress(L"NtGdiPolyPolyDraw");
	if (m_KernelGDIFunction.NtGdiPolyPolyDraw == NULL || MmIsAddressValid(m_KernelGDIFunction.NtGdiPolyPolyDraw) == FALSE) {
		DPRINT("GDI.cpp NtGdiPolyPolyDraw() = NULL An Error.KernelGDIFunctionInitialization Internal Function\r\n");
		goto $EXIT;
	}
	

	m_KernelGDIFunction.NtGdiSetPixel = (COLORREF(__cdecl *)(HDC, INT, INT, COLORREF))m_ShadowSSDT.Get3SDTFunAddress(L"NtGdiSetPixel");
	if (m_KernelGDIFunction.NtGdiSetPixel == NULL || MmIsAddressValid(m_KernelGDIFunction.NtGdiSetPixel) == FALSE) {
		DPRINT("GDI.cpp NtGdiSetPixel() = NULL An Error.KernelGDIFunctionInitialization Internal Function\r\n");
		goto $EXIT;
	}


	Status = STATUS_SUCCESS;
$EXIT:
	return Status;
}

PVOID GDI::AllocateVirtualMemory(SIZE_T Size)
{
	PVOID pMem = NULL;
	NtAllocateVirtualMemory(NtCurrentProcess(), &pMem, 0, &Size, MEM_COMMIT, PAGE_READWRITE);
	return pMem;
}

VOID GDI::FreeVirtualMemory(PVOID VirtualAddress, SIZE_T Size)
{
if (MmIsAddressValid(VirtualAddress))
	{
		NTSTATUS Status = NtFreeVirtualMemory(NtCurrentProcess(), &VirtualAddress, &Size, MEM_RELEASE);

		if (!NT_SUCCESS(Status)) {
			DPRINT("GDI.cpp Warning : Released memory failed.FreeVirtualMemory Internal Function\r\n");
		}
       return;
	
	}
	DPRINT("GDI.cpp Warning: Released memory does not exist.FreeVirtualMemory Internal Function\r\n");



}

HWND GDI::FindWindow(WCHAR * lpClassName, WCHAR * lpWindowName)
{
	SIZE_T allocSize = sizeof(UNICODE_STRING) * 2 + 260 * sizeof(CHAR) * 2;

	PVOID pMem = AllocateVirtualMemory(allocSize);

	HWND Hwnd = 0;

	if (pMem == NULL) {
		DPRINT("GDI.cpp AllocateVirtualMemory() An Error.FindWindow Internal Function\r\n");
		return 0;
	}

	RtlZeroMemory(pMem, allocSize);
	PUCHAR pBuffer = (PUCHAR)pMem;
	PUNICODE_STRING usClassName = (PUNICODE_STRING)(pBuffer);
	PUNICODE_STRING usWindowName = (PUNICODE_STRING)(pBuffer + sizeof(UNICODE_STRING));
	ULONG64 str_buffer = (ULONG64)pBuffer + sizeof(UNICODE_STRING) * 2;

	if (lpClassName)
	{
		RtlCopyMemory((PVOID)str_buffer, lpClassName, sizeof(WCHAR)*wcslen(lpClassName));
		RtlInitUnicodeString(usClassName, (PCWSTR)str_buffer);
		str_buffer += sizeof(WCHAR)*(wcslen(lpClassName) + 1);
	}
	if (lpWindowName)
	{
		RtlCopyMemory((PVOID)str_buffer, lpWindowName, sizeof(WCHAR)*wcslen(lpWindowName));
		RtlInitUnicodeString(usWindowName, (PCWSTR)str_buffer);
	}

	if (m_KernelGDIFunction.NtUserFindWindowEx != NULL) {
		Hwnd = m_KernelGDIFunction.NtUserFindWindowEx(HWND(0), HWND(0), usClassName, usWindowName, DWORD(0));
	}
	else {
		DPRINT("GDI.cpp NtUserFindWindowEx = NULL An Error.FindWindow Internal Function\r\n");
	}
	if (pMem != NULL) {
		FreeVirtualMemory(pMem, allocSize);
	}

	return Hwnd;
}

HDC GDI::GetDc(IN HWND hwnd)
{
	if (m_KernelGDIFunction.NtUserGetDC != NULL)
	{
    	return m_KernelGDIFunction.NtUserGetDC(hwnd);
	}

	DPRINT("GDI.cpp NtUserGetDC = NULL An Error.GetDc Internal Function\r\n");
	return NULL;
}


HWND GDI::FindWindowEx(HWND hWndParent, HWND hWndChildAfter, LPCWSTR lpszClass, LPCWSTR lpszWindow)
{
	SIZE_T allocSize = sizeof(UNICODE_STRING) * 2 + MAX_PATH * sizeof(WCHAR) * 2;
	auto pMem = AllocateVirtualMemory(allocSize);
	HWND hwnd = 0;
	if (pMem == NULL)
	{
		DPRINT("GDI.cpp AllocateVirtualMemory An Error.FindWindowEx Internal Function\r\n");
		return HWND(0);
	}

	RtlZeroMemory(pMem, allocSize);

	if (allocSize < sizeof(UNICODE_STRING) * 2 + MAX_PATH * sizeof(WCHAR) * 2)
	{
		DPRINT("GDI.cpp allocSize < sizeof(UNICODE_STRING) * 2 + MAX_PATH * sizeof(WCHAR) * 2.FindWindowEx Internal Function\r\n");
		return HWND(0);
	}

	auto pBuffer = (PUCHAR)pMem;
	PUNICODE_STRING usClassName = (PUNICODE_STRING)(pBuffer);
	PUNICODE_STRING usWindowName = (PUNICODE_STRING)(pBuffer + sizeof(UNICODE_STRING));
	auto str_buffer = pBuffer + sizeof(UNICODE_STRING) * 2;

	if (lpszClass)
	{
		RtlCopyMemory(str_buffer, lpszClass, sizeof(WCHAR)*wcslen(lpszClass));
		RtlInitUnicodeString(usClassName, (PCWSTR)str_buffer);
		str_buffer += sizeof(WCHAR)*(wcslen(lpszClass) + 1);
	}
	if (lpszWindow)
	{
		RtlCopyMemory(str_buffer, lpszWindow, sizeof(WCHAR)*wcslen(lpszWindow));
		RtlInitUnicodeString(usWindowName, (PCWSTR)str_buffer);
	}

	if (m_KernelGDIFunction.NtUserFindWindowEx != NULL) {
		hwnd = m_KernelGDIFunction.NtUserFindWindowEx(hWndParent, hWndChildAfter, usClassName, usWindowName, DWORD(0));
	}
	else {
		DPRINT("GDI.cpp Line NtUserFindWindowEx = NULL Triggers An Error.FindWindowEx Internal Function\r\n");
	}
	if (pMem != NULL) {
		FreeVirtualMemory(pMem, allocSize);
	}


	return hwnd;
}


COLORREF GDI::SetTextColor(HDC hdc, COLORREF crColor)
{
	PDC_ATTR pdcattr;
	COLORREF crOldColor;
	pdcattr = GdiGetDcAttr(hdc);
	if (pdcattr == NULL)
	{
		return CLR_INVALID;
	}

	crOldColor = (COLORREF)pdcattr->ulForegroundClr;
	pdcattr->ulForegroundClr = (ULONG)crColor;

	if (pdcattr->crForegroundClr != crColor)
	{
		pdcattr->ulDirty_ |= (DIRTY_TEXT | DIRTY_LINE | DIRTY_FILL);
		pdcattr->crForegroundClr = crColor;
	}

	return crOldColor;
}



BOOL GDI::ExtTextOutW(HDC hdc, INT x, INT y, UINT fuOptions, RECT * lprc, LPWSTR lpString, UINT cwc, INT * lpDx)
{
	BOOL		nRet = FALSE;
	PVOID       local_lpString = NULL;
	RECT*       local_lprc = NULL;
	INT*        local_lpDx = NULL;

	if (lprc != NULL) 
	{
		SIZE_T Len = sizeof(RECT);
		local_lprc = (RECT *)AllocateVirtualMemory(Len);
		if (local_lprc != NULL)
		{
			__try
			{
				RtlZeroMemory(local_lprc, Len);
				RtlCopyMemory(local_lprc, lprc, Len);
			}
			__except (1)
			{
				DPRINT("GDI.cpp Line RtlCopyMemory  Triggers An Error.ExtTextOutW Internal Function\r\n");
				goto $EXIT;
			}
		}
		else 
		{
			DPRINT("GDI.cpp Line local_lprc = null  Triggers An Error.ExtTextOutW Internal Function\r\n");
			goto $EXIT;
		}

	}

	if (cwc != 0)
	{
		SIZE_T     AllocSize = sizeof(WCHAR)*cwc + 1;
		local_lpString = AllocateVirtualMemory(AllocSize);
		if (local_lpString != NULL)
		{
			__try
			{
				RtlZeroMemory(local_lpString, AllocSize);
				RtlCopyMemory(local_lpString, lpString, AllocSize);
			}
			__except (1)
			{
				DPRINT("GDI.cpp Line RtlCopyMemory  Triggers An Error.ExtTextOutW Internal Function\r\n");
				goto $EXIT;
			}
		}
		else
		{
			DPRINT("GDI.cpp Line local_lpString = null  Triggers An Error.ExtTextOutW Internal Function\r\n");
			goto $EXIT;
		}
	}


	if (local_lpDx != NULL)
	{
		SIZE_T     AllocSize = sizeof(INT);
		local_lpDx = (INT *)AllocateVirtualMemory(AllocSize);
		if (local_lpDx != NULL)
		{
			__try
			{
				RtlZeroMemory(local_lpString, AllocSize);
				*local_lpDx = *lpDx;
			}
			__except (1)
			{
				DPRINT("GDI.cpp Line RtlCopyMemory  Triggers An Error.ExtTextOutW Internal Function\r\n");
				goto $EXIT;
			}
		}
		else
		{
			DPRINT("GDI.cpp Line local_lpDx = null  Triggers An Error.ExtTextOutW Internal Function\r\n");
		}
	}
	

	if (m_KernelGDIFunction.NtGdiExtTextOutW != NULL) {
		nRet = m_KernelGDIFunction.NtGdiExtTextOutW(hdc, x, y, fuOptions, local_lprc, (LPWSTR)local_lpString, cwc, local_lpDx, 0);
	}
	else {
		DPRINT("GDI.cpp Line NtGdiExtTextOutW = NULL Triggers An Error.TextOutW Internal Function\r\n");
	}
$EXIT:
	if (lprc != NULL)
	{
		FreeVirtualMemory(lprc, sizeof(RECT));
		lprc = NULL;
	}

	if (local_lpDx != NULL)
	{
		FreeVirtualMemory(local_lpDx, sizeof(INT));
		local_lpDx = NULL;
	}

	if (local_lpString != NULL)
	{
		FreeVirtualMemory(local_lpString, cwc);
		local_lpString = NULL;
	}

	return nRet;
}


BOOL GDI::ExtTextOutA(HDC hdc, INT x, INT y, UINT fuOptions,  RECT * lprc, LPCSTR lpString, UINT cch,  INT * lpDx)
{
	ANSI_STRING StringA;
	UNICODE_STRING StringU;
	BOOL ret;
	RtlInitAnsiString(&StringA, (LPSTR)lpString);
	RtlAnsiStringToUnicodeString(&StringU, &StringA, TRUE);
	ret = ExtTextOutW(hdc, x, y, fuOptions, lprc, StringU.Buffer, cch, lpDx);
	RtlFreeUnicodeString(&StringU);
	return ret;
}


ULONG_PTR GDI::Polyline(HDC hdc, const POINT * apt, INT cpt)
{


	SIZE_T allocSize = sizeof(DWORD) + cpt * sizeof(POINT);
	SIZE_T okSize = allocSize;
	ULONG_PTR nRet = 0;
	auto pMem = AllocateVirtualMemory(allocSize);
	if (!pMem)
	{
		DPRINT("GDI.cpp pMem == NULL Triggers An Error.Polyline Internal Function\r\n");
		return FALSE;
	}
	
	RtlZeroMemory(pMem, allocSize);
	if (allocSize < okSize)
	{
		return FALSE;
	}
	RtlCopyMemory(pMem, apt, sizeof(POINT)*cpt);
	auto lpcpt = (PDWORD)((PUCHAR)pMem + sizeof(POINT)*cpt);
	*lpcpt = cpt;

	if (m_KernelGDIFunction.NtGdiPolyPolyDraw != NULL) {
		nRet = m_KernelGDIFunction.NtGdiPolyPolyDraw(hdc, (POINT *)pMem, (ULONG*)lpcpt, INT64(1), 4);
		DPRINT("%d \r\n", nRet);
	}
	else {
		DPRINT("GDI.cpp NtGdiPolyPolyDraw == NULL Triggers An Error.Polyline Internal Function\r\n");
	}

	if (pMem != NULL) {
		FreeVirtualMemory(pMem, allocSize);
	}

	return nRet;

}

COLORREF GDI::SetPixel(HDC hdc, INT x, INT y, COLORREF crColor)
{

	COLORREF nRet;
	if (m_KernelGDIFunction.NtGdiSetPixel != NULL) {
		nRet = m_KernelGDIFunction.NtGdiSetPixel(hdc, x, y, crColor);
	}
	else {
		DPRINT("GDI.cpp NtGdiSetPixel == NULL Triggers An Error.SetPixel Internal Function\r\n");
	}

	return nRet;
}
